package com.yuyu.tt.exec;

import lombok.AllArgsConstructor;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.AsyncContext;
import javax.servlet.http.HttpServletRequest;

/**
 * @author 深海鱼鱼
 * @version 1.0
 * @description:
 * @date 2024/11/20 15:45
 */
@Component
public class InheritableRequestTaskExecutor extends ThreadPoolTaskExecutor {
    private void initTaskExec() {
        setCorePoolSize(20);
        setMaxPoolSize(100);
        setQueueCapacity(200);
        setKeepAliveSeconds(60);
        setAllowCoreThreadTimeOut(true);
        setThreadNamePrefix("TaskExec-");
    }

    @Override
    public void execute(Runnable task) {
        initTaskExec();
        // 获取当前线程的请求上下文
        RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
        RequestContextHolder.setRequestAttributes(requestAttributes, true);
        AsyncContext asyncContext = null;
        if (requestAttributes != null) {
            HttpServletRequest request = ((ServletRequestAttributes) requestAttributes).getRequest();
            asyncContext = request.startAsync();
        }
        AsyncContext finalAsyncContext = asyncContext;
        super.execute(new RequestContextRunnable(task, finalAsyncContext));
    }

    @AllArgsConstructor
    private static class RequestContextRunnable implements Runnable {
        private final Runnable task;
        private final AsyncContext finalAsyncContext;

        @Override
        public void run() {
            try {
                task.run();
            } finally {
                // 执行完成后清理请求上下文
                if (finalAsyncContext != null) {
                    finalAsyncContext.complete();
                }
            }
        }
    }
}
